package org.zkdn.harborimagedownload.controller;

import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONArray;
import com.alibaba.fastjson2.JSONObject;
import lombok.AllArgsConstructor;
import lombok.Builder;
import lombok.Data;
import lombok.NoArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.compress.archivers.tar.TarArchiveEntry;
import org.apache.commons.compress.archivers.tar.TarArchiveOutputStream;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.*;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.client.RestTemplate;
import org.springframework.web.reactive.function.client.WebClient;
import org.zkdn.harborimagedownload.config.HidConfig;
import reactor.core.publisher.Mono;

import java.io.BufferedOutputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.util.*;
import java.util.concurrent.CompletableFuture;

@Slf4j
@RestController
@RequestMapping("/harbor")
public class HarborImageController {

    @Autowired
    private RestTemplate restTemplate;

//    @Autowired
    private WebClient webClient;

    @Autowired
    private HidConfig hidConfig;

    public HttpEntity<Object> getRequestEntity() {
        HttpHeaders httpHeaders = new HttpHeaders();
        httpHeaders.setBasicAuth(hidConfig.getHarbor().getUsername(), hidConfig.getHarbor().getPassword());
        return new HttpEntity<>(httpHeaders);
    }

    @GetMapping("/list")
    public Mono<JSONArray> list() {
        String baseUrl = hidConfig.getHarbor().getUrl() + "/api/v2.0/projects";
        ResponseEntity<JSONArray> response = restTemplate.exchange(baseUrl, HttpMethod.GET, getRequestEntity(), JSONArray.class);
        JSONArray projects = response.getBody();
        log.info(projects.toJSONString());
        return Mono.just(projects);
    }

    @GetMapping("/projects/list")
    public Mono<?> projects() {
        ResponseEntity<JSONArray> response = restTemplate.exchange(hidConfig.getHarbor().getUrl() + "/api/v2.0/projects", HttpMethod.GET, getRequestEntity(), JSONArray.class);
        JSONArray projects = response.getBody();
        log.info(projects.toJSONString());
        return Mono.just(projects);
    }

    @GetMapping("/images/list")
    public Mono<?> images(String projectName) {
        ResponseEntity<JSONArray> repository = restTemplate.exchange(hidConfig.getHarbor().getUrl() + "/api/v2.0/projects/" + projectName + "/repositories?page_size=100&page=1", HttpMethod.GET, getRequestEntity(), JSONArray.class);
        JSONArray body = repository.getBody();
        return Mono.just(body);
    }

    @GetMapping("/artifacts/list")
    public Mono<?> artifacts(String projectName, String imageName) {
        String url = hidConfig.getHarbor().getUrl() + "/api/v2.0/projects/" + projectName + "/repositories/" + imageName + "/artifacts?page_size=100&page=1";
        ResponseEntity<JSONArray> artifacts = restTemplate.exchange(url, HttpMethod.GET, getRequestEntity(), JSONArray.class);
        return Mono.just(artifacts.getBody());
    }

    @GetMapping("/tags/list")
    public Mono<?> tags(String projectName, String imageName, String artifactName) {
        String url = hidConfig.getHarbor().getUrl() + "/api/v2.0/projects/" + projectName + "/repositories/" + imageName + "/artifacts/" + artifactName + "/tags";
        ResponseEntity<JSONArray> tags = restTemplate.exchange(url, HttpMethod.GET, getRequestEntity(), JSONArray.class);
        return Mono.just(tags.getBody());
    }

    @GetMapping(value = "/download/image", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Mono<ResponseEntity<byte[]>> download(String projectName, String imageName, String tagName) {
        return Mono.fromCallable(() -> {
            try (ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
                 BufferedOutputStream tarBufferedOutput = new BufferedOutputStream(byteArrayOutputStream);
                 TarArchiveOutputStream tarArchiveOutputStream = new TarArchiveOutputStream(tarBufferedOutput);
                 ) {
                String url = hidConfig.getHarbor().getUrl() + "/v2/" + projectName + "/" + imageName + "/manifests/" + tagName;
                ResponseEntity<JSONObject> manifests = restTemplate.exchange(url, HttpMethod.GET, getRequestEntity(), JSONObject.class);
                JSONObject body = manifests.getBody();
                log.info("获取镜像信息 {}", body.toJSONString());

                buildConfigJson(projectName, imageName, body, tarArchiveOutputStream);

                ArrayList<Object> manifest = new ArrayList<>();
                HashMap<String, Object> manifestObject = new HashMap<>();
                manifestObject.put("Config", body.getJSONObject("config").getString("digest").replace("sha256:", "") + ".json");
                manifestObject.put("RepoTags", Arrays.asList(projectName + "/" + imageName + ":" + tagName));

                JSONArray layers = body.getJSONArray("layers");
                writeRepositoriesFile(tarArchiveOutputStream, projectName, imageName, tagName, layers.getJSONObject(layers.size() - 1).getString("digest").replace("sha256:", ""));

                ArrayList<String> mylayers = new ArrayList<>();
                long start = System.currentTimeMillis();
                CompletableFuture[] allSubFeature = layers.parallelStream().map(layer -> CompletableFuture.supplyAsync(() -> {
                    String digest = ((LinkedHashMap) layer).get("digest").toString();
                    String url1 = hidConfig.getHarbor().getUrl() + "/v2/" + projectName + "/" + imageName + "/blobs/" + digest;
                    byte[] layerData = restTemplate.exchange(url1, HttpMethod.GET, getRequestEntity(), byte[].class).getBody();
                    String layerName = digest.replace("sha256:", "") + ".tar";
                    log.info("下载层 layer {}", digest);
                    return LayerEntity.builder().layer(layerName).data(layerData).build();
                })).toArray(CompletableFuture[]::new);
                CompletableFuture.allOf(allSubFeature).thenApply(v -> {
                    for (CompletableFuture<LayerEntity> future : allSubFeature) {
                        try {
                            LayerEntity entity = future.get();
                            TarArchiveEntry entry = new TarArchiveEntry(entity.getLayer());
                            entry.setSize(entity.getData().length);
                            tarArchiveOutputStream.putArchiveEntry(entry);
                            tarArchiveOutputStream.write(entity.getData());
                            tarArchiveOutputStream.closeArchiveEntry();
                        } catch (Exception e) {
                            log.error(e.getMessage(), e);
                        }
                    }
                    return null;
                }).get();
                log.info("下载结束 {} ms", System.currentTimeMillis() - start);
                manifestObject.put("Layers", mylayers);
                manifest.add(manifestObject);

                buildMainfestJson(manifest, tarArchiveOutputStream);

                tarArchiveOutputStream.flush();
                return Mono.just(
                        ResponseEntity.ok()
                                .header(HttpHeaders.CONTENT_DISPOSITION, "attachment; filename=\"" + projectName + "-" + imageName + ":" + tagName + ".tar\"")
                                .body(byteArrayOutputStream.toByteArray())
                );
            } catch (Exception e) {
                log.error(e.getMessage(), e);
                return Mono.just(
                        ResponseEntity.status(500).body(e.getMessage().getBytes())
                );
            }
        }).flatMap(response -> response);
    }

    private static void buildMainfestJson(ArrayList<Object> manifest, TarArchiveOutputStream tarArchiveOutputStream) throws IOException {
        String manifestStr = JSON.toJSONString(manifest);
        log.info("manifest {}", manifestStr);
        byte[] manifestbytes = manifestStr.getBytes("UTF-8");
        TarArchiveEntry repositoriesEntry = new TarArchiveEntry("manifest.json");
        repositoriesEntry.setSize(manifestbytes.length);
        tarArchiveOutputStream.putArchiveEntry(repositoriesEntry);
        tarArchiveOutputStream.write(manifestbytes);
        tarArchiveOutputStream.closeArchiveEntry();
    }

    private void buildConfigJson(String projectName, String imageName, JSONObject body, TarArchiveOutputStream tarArchiveOutputStream) throws IOException {
        String configDigest = body.getJSONObject("config").getString("digest");
        String url = hidConfig.getHarbor().getUrl() + "/v2/" + projectName + "/" + imageName + "/blobs/" + configDigest;
        String config = restTemplate.exchange(url, HttpMethod.GET, getRequestEntity(), String.class).getBody();
        String configDigestJson = configDigest.replace("sha256:", "") + ".json";
        log.info("configDigestJson {}", config);
        byte[] repositoriesData = config.getBytes("UTF-8");
        TarArchiveEntry repositoriesEntry = new TarArchiveEntry(configDigestJson);
        repositoriesEntry.setSize(repositoriesData.length);
        tarArchiveOutputStream.putArchiveEntry(repositoriesEntry);
        tarArchiveOutputStream.write(repositoriesData);
        tarArchiveOutputStream.closeArchiveEntry();
    }

    private static void writeRepositoriesFile(TarArchiveOutputStream tarOutput, String repository, String imageName, String tag, String digest) throws IOException {
        // 创建 repositories 文件内容
        Map<String, Map<String, String>> repositories = new HashMap<>();
        Map<String, String> tags = new HashMap<>();
        tags.put(tag, digest);  // 确保 tags 中使用最新的层的哈希值
        repositories.put(repository + "/" + imageName, tags);

        JSONObject repositoriesJson = new JSONObject(repositories);
        log.info("repositories {}", repositoriesJson.toJSONString());

        // 将 repositories 文件内容写入到 tar 中
        byte[] repositoriesData = repositoriesJson.toJSONString().getBytes("UTF-8");
        TarArchiveEntry repositoriesEntry = new TarArchiveEntry("repositories");
        repositoriesEntry.setSize(repositoriesData.length);
        tarOutput.putArchiveEntry(repositoriesEntry);
        tarOutput.write(repositoriesData);
        tarOutput.closeArchiveEntry();

    }

    @Data
    @Builder
    @AllArgsConstructor
    @NoArgsConstructor
    public static class LayerEntity {
        private String layer;
        private byte[] data;
    }

}
